import numpy as np
from scipy.signal import welch, find_peaks
from scipy.interpolate import interp1d
from scipy.stats import iqr
from scipy.linalg import norm
from sklearn.decomposition import PCA
from numpy.linalg import norm
import numpy as np
from scipy.signal import butter, filtfilt
from scipy.signal import coherence, csd, get_window

def step1(q1, q2, q3, q4, q5, q6):
    fl = 3.0
    fh = 12.0
    fs = 100.0
    bw = 2
    nFFT = 128
    nOverlap = nFFT // 2

    nameTable = []
    fea = []
    print("Received q1 type:", type(q1))
    qE2W = np.array(q1)
    qE2T = np.array(q2)
    qE2F = np.array(q3)
    qE2M = np.array(q4)
    qE2R = np.array(q5)
    qE2L = np.array(q6)
    print("qE2W shape:",qE2W.shape)
    print()
    print(f"qE2W: {qE2W[0, 0]}, "
          f"qE2T: {qE2T[0, 0]}, "
          f"qE2F: {qE2F[0, 0]}, "
          f"qE2M: {qE2M[0, 0]}, "
          f"qE2R: {qE2R[0, 0]}, "
          f"qE2L: {qE2L[0, 0]}")
    # wrist to others: qW2X = inv(qE2W) * qE2X
    qW2T = quat_multiply(quat_inv(qE2W), qE2T)
    qW2F = quat_multiply(quat_inv(qE2W), qE2F)
    qW2M = quat_multiply(quat_inv(qE2W), qE2M)
    qW2R = quat_multiply(quat_inv(qE2W), qE2R)
    qW2L = quat_multiply(quat_inv(qE2W), qE2L)

    qCell = [qE2W, qE2T, qE2F, qE2M, qE2R, qE2L]
    qR = [qW2T, qW2F, qW2M, qW2R, qW2L]

    dataLength = qW2L.shape[0]
    angP = np.zeros((dataLength, 6))
    angO = np.zeros((3, 6))

    # For each sensor compute Euler angles, unbias and extract principal components
    for i in range(6):
        tempEAs = getEulerAngle(qCell[i], 'zyx')  # shape Nx3
        angP[:, i], angO[:, i], _ = unBiasAndPre(tempEAs)

    print(angP[0,0])


    #return flip_columns_except_second(angP), flip_columns_except_second(angO)
    #return flip_columns_except_05(angP), flip_columns_except_05(angO)
    for i in range(angO.shape[1]):
        # 找到该主成分方向中最大绝对值元素的索引
        idx = np.argmax(np.abs(angO[:, i]))
        # 若该元素为负，则整体反向
        if angO[idx, i] < 0:
            angO[:, i] = -angO[:, i]
            angP[:, i] = -angP[:, i]
    return angP, angO

def step2(v, nFFT, nOverlap, fl, fh, fs, bw):
    v = np.array(v)
    f, Pxx = welch(v, window='hann', nperseg=nFFT, noverlap=nOverlap, fs=fs)
    mask = (f >= fl) & (f <= fh)
    tempPxx = Pxx[mask]
    tempF = f[mask]
    idx = np.argmax(tempPxx)
    f0 = tempF[idx]
    return f0


# Supporting functions (some need further implementation details from MATLAB versions):
def step3(vF, nFFT, nOverlap, fs, fl, fh):
    vF = np.array(vF)
    F0, F50, SF50, K = alPSD(vF, nFFT, nOverlap, fs, fl, fh)
    return F0, F50, SF50, K

def step4(F,F0,F50,SF50,K,angO,q1, q2, q3, q4, q5, q6,addname):
    nameTable = []
    fea = []
    addName = addname
    fl = 3.0
    fh = 12.0
    fs = 100.0
    bw = 2
    nFFT = 128
    nOverlap = nFFT // 2
    angF = np.array(F)
    angF0 = np.array(F0)
    angF50 = np.array(F50)
    angSF50 = np.array(SF50)
    angK = np.array(K)
    angO = np.array(angO)

    print("Received q1 type:", type(q1))
    qE2W = np.array(q1)
    qE2T = np.array(q2)
    qE2F = np.array(q3)
    qE2M = np.array(q4)
    qE2R = np.array(q5)
    qE2L = np.array(q6)
    print("qE2W shape:",qE2W.shape,qE2W[0,0],qE2T[0,0],qE2F[0,0],qE2M[0,0],qE2R[0,0],qE2L[0,0])
    # wrist to others: qW2X = inv(qE2W) * qE2X
    qW2T = quat_multiply(quat_inv(qE2W), qE2T)
    qW2F = quat_multiply(quat_inv(qE2W), qE2F)
    qW2M = quat_multiply(quat_inv(qE2W), qE2M)
    qW2R = quat_multiply(quat_inv(qE2W), qE2R)
    qW2L = quat_multiply(quat_inv(qE2W), qE2L)

    qCell = [qE2W, qE2T, qE2F, qE2M, qE2R, qE2L]
    qR = [qW2T, qW2F, qW2M, qW2R, qW2L]


    # Dynamic stats for each sensor
    angPeakAvg = np.zeros(6)
    angPeakCV = np.zeros(6)
    angPeakMax = np.zeros(6)
    angFreAvg = np.zeros(6)
    angFreCV = np.zeros(6)
    angFreIQR = np.zeros(6)
    angPeakIQR = np.zeros(6)
    angRotZ = np.zeros(6)




    for i in range(6):
        angf, angdf, angm, angdm = dynamicAnalyze(angF[:, i], fs)
        angPeakAvg[i] = np.mean(angm)
        angPeakCV[i] = np.std(angm, ddof=1) / np.mean(angm)
        angPeakMax[i] = np.max(angm)
        angFreAvg[i] = np.mean(angf)
        angFreCV[i] = np.std(angf, ddof=1) / np.mean(angf)
        angFreIQR[i] = iqr(angdf)
        angPeakIQR[i] = iqr(angdm) / np.sqrt(np.mean(angdm**2))
        angRotZ[i] = abs(angO[2, i])

    k = 0
    angPhase = np.zeros(15)
    AngR = np.zeros(5)

    for i in range(5):
        _, _, AngR[i] = alQua(qR[i])
        for j in range(i + 1, 6):
            tempFL = fl
            tempFH = fh
            _, angPhase[k] = alPhase(angF[:, i], angF[:, j], tempFL, tempFH, fs)
            k += 1

    # Compose feature names and append features
    names_phase = ['angPhaseW2T', 'angPhaseW2F', 'angPhaseW2M', 'angPhaseW2R', 'angPhaseW2L',
                   'angPhaseT2F', 'angPhaseT2M', 'angPhaseT2R', 'angPhaseT2L',
                   'angPhaseF2M', 'angPhaseF2R', 'angPhaseF2L',
                   'angPhaseM2R', 'angPhaseM2L', 'angPhaseR2L']
    names_angR = ['AngRT2W', 'AngRT2F', 'AngRT2M', 'AngRT2R', 'AngRT2L']

    nameTable.extend([name + addName for name in (names_phase + names_angR)])
    fea.extend(angPhase.tolist() + AngR.tolist())

    assert len(nameTable) == len(fea)

    # RMS values
    angRMS = np.array([rms(angF[:, i]) for i in range(6)])

    # Feature names for other features
    names_rest = ['angF0_W', 'angF0_T', 'angF0_F', 'angF0_M', 'angF0_R', 'angF0_L',
                  'angF50_W', 'angF50_T', 'angF50_F', 'angF50_M', 'angF50_R', 'angF50_L',
                  'angSF50_W', 'angSF50_T', 'angSF50_F', 'angSF50_M', 'angSF50_R', 'angSF50_L',
                  'angK_W', 'angK_T', 'angK_F', 'angK_M', 'angK_R', 'angK_L',
                  'angPeakAvg_W', 'angPeakAvg_T', 'angPeakAvg_F', 'angPeakAvg_M', 'angPeakAvg_R', 'angPeakAvg_L',
                  'angPeakCV_W', 'angPeakCV_T', 'angPeakCV_F', 'angPeakCV_M', 'angPeakCV_R', 'angPeakCV_L',
                  'angPeakMax_W', 'angPeakMax_T', 'angPeakMax_F', 'angPeakMax_M', 'angPeakMax_R', 'angPeakMax_L',
                  'angPeakIQR_W', 'angPeakIQR_T', 'angPeakIQR_F', 'angPeakIQR_M', 'angPeakIQR_R', 'angPeakIQR_L',
                  'angFreAvg_W', 'angFreAvg_T', 'angFreAvg_F', 'angFreAvg_M', 'angFreAvg_R', 'angFreAvg_L',
                  'angFreCV_W', 'angFreCV_T', 'angFreCV_F', 'angFreCV_M', 'angFreCV_R', 'angFreCV_L',
                  'angFreIQR_W', 'angFreIQR_T', 'angFreIQR_F', 'angFreIQR_M', 'angFreIQR_R', 'angFreIQR_L',
                  'angRMS_W', 'angRMS_T', 'angRMS_F', 'angRMS_M', 'angRMS_R', 'angRMS_L']

    nameTable.extend([name + addName for name in names_rest])

    fea.extend(angF0.tolist())
    fea.extend(angF50.tolist())
    fea.extend(angSF50.tolist())
    fea.extend(angK.tolist())
    fea.extend(angPeakAvg.tolist())
    fea.extend(angPeakCV.tolist())
    fea.extend(angPeakMax.tolist())
    fea.extend(angPeakIQR.tolist())
    fea.extend(angFreAvg.tolist())
    fea.extend(angFreCV.tolist())
    fea.extend(angFreIQR.tolist())
    fea.extend(angRMS.tolist())

    assert len(nameTable) == len(fea)

    # Finding max tremor index
    peakAvgNorm = angPeakAvg / np.sqrt(np.sum(angPeakAvg**2))
    angRMSNorm = angRMS / np.sqrt(np.sum(angRMS**2))
    tremorNorm = peakAvgNorm + angRMSNorm
    maxIdxTremor = np.argmax(tremorNorm)  # zero-based index in Python

    maxIdx = maxIdxTremor + 1  # convert to 1-based to match MATLAB indexing if needed

    nameTable.append('maxIdx' + addName)
    fea.append(maxIdx)

    nameTable.extend([f + addName for f in ['angF0', 'angF50', 'angSF50', 'angK',
                                            'angPeakAvg', 'angPeakCV', 'angPeakMax', 'angPeakIQR',
                                            'angFreAvg', 'angFreCV', 'angFreIQR', 'angRotZ', 'angRMS']])
    fea.extend([angF0[maxIdxTremor], angF50[maxIdxTremor], angSF50[maxIdxTremor], angK[maxIdxTremor],
                angPeakAvg[maxIdxTremor], angPeakCV[maxIdxTremor], angPeakMax[maxIdxTremor], angPeakIQR[maxIdxTremor],
                angFreAvg[maxIdxTremor], angFreCV[maxIdxTremor], angFreIQR[maxIdxTremor], angRotZ[maxIdxTremor], angRMS[maxIdxTremor]])

    assert len(nameTable) == len(fea)
    fea = np.array(fea).T
    return nameTable, fea

def quat_inv(q):
    # Quaternion inverse assuming unit quaternion
    q_conj = q.copy()
    q_conj[:, 1:4] *= -1
    norm_q = np.sum(q * q, axis=1, keepdims=True)
    return q_conj / norm_q

def quat_multiply(q, r):
    # Quaternion multiply q * r for arrays Nx4; q,r shape (N,4)
    w0, x0, y0, z0 = q[:,0], q[:,1], q[:,2], q[:,3]
    w1, x1, y1, z1 = r[:,0], r[:,1], r[:,2], r[:,3]
    return np.column_stack([
        w0*w1 - x0*x1 - y0*y1 - z0*z1,
        w0*x1 + x0*w1 + y0*z1 - z0*y1,
        w0*y1 - x0*z1 + y0*w1 + z0*x1,
        w0*z1 + x0*y1 - y0*x1 + z0*w1])

def getEulerAngle(quat, order='zyx'):
    angle_z = np.arctan2(2 * (quat[:,1] * quat[:,2] + quat[:,0] * quat[:,3]), 1 - 2 * (quat[:,2] * quat[:,2] + quat[:,3] * quat[:,3]))
    angle_y = -np.arcsin(np.clip(2 * (quat[:,1] * quat[:,3] - quat[:,0] * quat[:,2]), -1.0, 1.0))
    angle_x = np.arctan2(2 * (quat[:,0] * quat[:,1] + quat[:,2] * quat[:,3]), 1 - 2 * (quat[:,1] * quat[:,1] + quat[:,2] * quat[:,2]))
    # unwrap angles to prevent jumps
    angles = np.stack((angle_z, angle_y, angle_x), axis=1)
    angles = np.unwrap(angles, axis=0)
    return np.rad2deg(angles)

def unBiasAndPre(v):
    # Remove bias by wavelet filter (not implemented here, use simple high-pass or demean)
    v_filtered = np.apply_along_axis(lambda x: filterByWavelet(x, 5), 0, v)
    pca = PCA(n_components=3)
    score = pca.fit_transform(v_filtered)
    coeff = pca.components_.T
    latent = pca.explained_variance_
    vP = score[:, 0]
    oV = coeff[:, 0]
    lat = latent[0] / np.sum(latent)
    return vP, oV, lat

def flip_columns_except_second(matrix: np.ndarray) -> np.ndarray:
    """
    保持第2列符号不变，其余列取反
    matrix: numpy二维数组
    """
    result = matrix.copy()
    for j in range(matrix.shape[1]):
        if j != 1:  # 不是第二列就取反
            result[:, j] = -result[:, j]
    return result

def flip_columns_except_05(matrix: np.ndarray) -> np.ndarray:
    """
    保持第2列符号不变，其余列取反
    matrix: numpy二维数组
    """
    result = matrix.copy()
    for j in range(matrix.shape[1]):
        if j != 0 and j!=5 :  # 不是第二列就取反
            result[:, j] = -result[:, j]
    return result

def filterByWavelet(signal, level):
    import pywt
    """
    信号去漂移滤波，类似MATLAB的wavedec + wrcoef组合
    参数：
        signal: 1维numpy数组，输入信号
        level: 小波分解层数，整数
    返回：
        filtered_signal: 去除低频漂移后的信号
    """
    wname = 'db4'
    # 多尺度小波分解
    coeffs = pywt.wavedec(signal, wavelet=wname, level=level)
    # 重构出近似系数（趋势部分）
    drift = pywt.waverec([coeffs[0]] + [np.zeros_like(c) for c in coeffs[1:]], wavelet=wname)
    drift = drift[:len(signal)]  # waverec可能导致长度略有不同，截断匹配
    # 去除漂移，得到滤波后信号
    filtered_signal = signal - drift
    return filtered_signal


def alPSD(w, nFFT, nOverlap, fs, fl, fh):
    f, Pxx = welch(w, window='hann', nperseg=nFFT, noverlap=nOverlap, fs=fs)
    df = f[1] - f[0]
    mask = (f >= fl) & (f <= fh)
    PIn = Pxx[mask]
    fIn = f[mask]
    idxMax = np.argmax(PIn)
    F0 = fIn[idxMax]
    K = 10 * np.log10(PIn[idxMax])
    NPS = np.cumsum(Pxx) / np.sum(Pxx)
    idx50 = np.where(NPS > 0.5)[0][0]
    if abs(NPS[idx50] - 0.5) > abs(NPS[idx50 - 1] - 0.5):
        idx50 = idx50 - 1
    F50 = f[idx50]

    sLength = idx50 - 1
    bLength = len(f) - idx50 -1
    tempLength = max(sLength, bLength)
    print(sLength, bLength, tempLength)
    pLessF50 = np.pad(Pxx[:idx50 - 1], (0, tempLength - sLength), 'constant')
    pLargerF50 = np.pad(Pxx[idx50 + 1:], (0, tempLength - bLength), 'constant')
    print(len(pLessF50), len(pLargerF50), tempLength)
    NSFP = np.cumsum(np.concatenate(([Pxx[idx50]], pLessF50 + pLargerF50))) / np.sum(Pxx)
    idxSF50 = np.where(NSFP > 0.58)[0][0]
    SF50 = df * (idxSF50)
    return F0, F50, SF50, K

def dynamicAnalyze(v, fs):
    fs1 = 2000
    t0 = np.arange(len(v)) / fs
    t1 = np.arange(0, t0[-1], 1/fs1)
    f_interp = interp1d(t0, v, kind='cubic')
    v_interp = f_interp(t1)
    th = 4 * fs1 / fs
    dv = np.diff(v_interp)
    mv = v_interp[1:] * v_interp[:-1]
    upCross = np.where((dv > 0) & (mv < 0))[0]
    upCrossRes = [upCross[0]]
    j = 1
    for i in range(1, len(upCross)):
        dUpCross = upCross[i] - upCross[i-1]
        if dUpCross < th:
            if i == len(upCross) - 1:
                upCrossRes[-1] = upCross[i]
            elif i == 1:
                pass
            else:
                temp = (upCross[i+1] + upCross[i-2]) / 2
                if abs(upCross[i] - temp) < abs(upCross[i-1] - temp):
                    upCrossRes[-1] = upCross[i]
        else:
            upCrossRes.append(upCross[i])
            j += 1
    upCross = np.array(upCrossRes)
    f = np.zeros(len(upCross) - 2)
    df = np.zeros(len(upCross) - 2)
    m = np.zeros(len(upCross) - 2)
    dm = np.zeros(len(upCross) - 2)
    for i in range(len(upCross) - 2):
        tempF = fs1 / (upCross[i+1] - upCross[i])
        tempFNext = fs1 / (upCross[i+2] - upCross[i+1])
        tempM = np.max(v_interp[upCross[i]+1:upCross[i+1]+1])
        tempMNext = np.max(v_interp[upCross[i+1]+1:upCross[i+2]+1])
        f[i] = tempF
        df[i] = -tempF + tempFNext
        m[i] = tempM
        dm[i] = -tempM + tempMNext
    return f, df, m, dm

def alQua(qua):
    # Implementation of quaternion-based analysis similar to MATLAB
    quaMiddle, middleIdx, minDist = findMiddleQua(qua)
    qua_inv_middle = quat_inv_single(quaMiddle)
    q_rel = quat_multiply_single(qua_inv_middle, qua).T
    ang = np.arctan2(np.sign(q_rel[:, 2]) * np.sqrt(np.sum(q_rel[:, 1:4]**2, axis=1)), q_rel[:, 0]) * 2 * 180 / np.pi
    rot = np.sign(q_rel[:, 2])[:, None] * q_rel[:, 1:4] / np.linalg.norm(q_rel[:, 1:4], axis=1)[:, None]
    tempRot = (rot[middleIdx - 1] + rot[middleIdx + 1]) / 2
    tempRot = tempRot / np.linalg.norm(tempRot)
    rot[middleIdx] = tempRot
    angFront = ang[:-1]
    angBehind = ang[1:]
    angDiff = np.diff(ang)
    tempangBehind = -angBehind
    angDiffNeg = tempangBehind - angFront
    tempIdx = np.where(np.abs(angDiffNeg) < np.abs(angDiff))[0]
    for idx in tempIdx:
        ang[idx + 1:] = -ang[idx + 1:]
        rot[idx + 1:, :] = -rot[idx + 1:, :]
    return ang, rot, minDist

def findMiddleQua(qua):
    n = qua.shape[0]
    dist = np.zeros(n)
    for i in range(n):
        base = qua[i, :]
        comp = np.vstack((qua[:i, :], qua[i+1:, :]))
        q = quat_inv_single(base)
        q = q.T
        q = q.reshape(1, 4)
        qS02S = quat_multiply(q, comp)
        temp = np.abs(np.arccos(qS02S[:,0]) * 2)
        dist[i] = np.sum(temp)
    middleIdx = np.argmin(dist)
    minDist = dist[middleIdx]
    quaMiddle = qua[middleIdx, :]
    return quaMiddle, middleIdx, minDist

# Helper quaternion operations for single quaternion or batches if needed
def quat_inv_single(q):
    # q shape (4,), unit quaternion inverse is conjugate
    q_conj = q.copy()
    q_conj[1:4] *= -1
    return q_conj / np.dot(q, q)

def quat_multiply_single(q, r):
    # q, r shapes (4,)
    w0, x0, y0, z0 = q
    w1, x1, y1, z1 = r.T
    return np.array([
        w0*w1 - x0*x1 - y0*y1 - z0*z1,
        w0*x1 + x0*w1 + y0*z1 - z0*y1,
        w0*y1 - x0*z1 + y0*w1 + z0*x1,
        w0*z1 + x0*y1 - y0*x1 + z0*w1,
        ])

def rms(x):
    return np.sqrt(np.mean(np.square(x)))

def alPhase(ang1, ang2, fl, fh, fs):
    window = get_window('blackman', 512)
    f, tempC = coherence(ang1, ang2, fs=fs, window=window, nperseg=512, noverlap=256)

    f_cpsd, tempP = csd(ang1, ang2, fs=fs, window=window, nperseg=512, noverlap=256)

    # 频率筛选
    mask = (f >= fl) & (f <= fh)
    tempC = tempC[mask]
    tempP = tempP[mask]

    # 相位 (单位：度)
    tempPAngle = np.angle(tempP) / np.pi * 180
    tempPAngle = -tempPAngle
    # 找最大相干值
    CMax = np.max(tempC)
    Phase_candidates = tempPAngle[tempC == CMax]

    # 取第一个
    Phase = Phase_candidates[0] if len(Phase_candidates) > 0 else np.nan

    return CMax, Phase